#pragma once

#include <sys/mman.h>
#include <unistd.h>

#include <cstddef>
#include <cstring>
#include <memory>
#include <unordered_map>

#include "base/common.h"

#include "util/algo/lock.h"

namespace hawking {
namespace indexlib {

class MemoryChunk {
public:
    MemoryChunk(size_t n_size = 0) : size_(n_size) {
        pos_ = holder_ = sizeof(MemoryChunk);
        if (size_ < pos_) {
            size_ = pos_;
        }
    }

    /**
     * malloc from chunk
     * @param n_bytes number of bytes to malloc
     */
    void *Allocate(size_t n_bytes) {
        if (size_ - pos_ < n_bytes) {
            return nullptr;
        }
        void *ptr = (void *)((char *)(this) + pos_);
        pos_ += n_bytes;
        return ptr;
    }

    void *Allocate(size_t n_bytes, size_t align) {
        auto ptr = (void *)((char *)(this) + pos_);
        auto space = size_ - pos_;
        auto aligned_ptr = std::align(align, n_bytes, ptr, space);
        if (aligned_ptr) {
            pos_ = (char *)aligned_ptr - (char *)(this) + n_bytes;
        }
        return aligned_ptr;
    }

    bool IsInChunk(const void *ptr) const {
        return ptr >= (void *)this && ptr < (void *)((char *)this + pos_);
    }

    /**
     * is there any space to allocate
     * @return true if no free space
     */
    bool IsEmpty() const { return pos_ >= size_; }

    /** get total size of chunk */
    size_t TotalBytes() const { return size_; }

    /** get current allocation position */
    size_t Pos() const { return pos_; }

    /** get free size */
    size_t FreeSize() const { return size_ - pos_; }

    /** get used size */
    size_t UsedSize() const { return pos_ - holder_; }

    /** reset chunk to free stat */
    void Reset() { pos_ = holder_; }

    /** get holder **/
    size_t HolderSize() const { return holder_; }

    /** clear trunk data **/
    void Clear() {
        if (size_ <= holder_) {
            return;
        }
        auto nDataBytes = size_ - holder_;
        auto p = (char *)(this) + holder_;
        std::memset(p, 0, nDataBytes);
    }

protected:
    size_t size_;   // size of the chunk
    size_t pos_;    // current allocation position
    size_t holder_; // size of this object
};

class ChainedMemoryChunk : public MemoryChunk {
public:
    ChainedMemoryChunk(size_t n_size) : MemoryChunk(n_size), next_(nullptr), prev_(nullptr) {
        pos_ = holder_ = sizeof(ChainedMemoryChunk);
        if (size_ < pos_) {
            size_ = pos_;
        }
    }

    /** get next chunk*/
    inline ChainedMemoryChunk *next() const { return next_; }
    inline ChainedMemoryChunk *&next() { return next_; }

    inline ChainedMemoryChunk *prev() const { return prev_; }
    inline ChainedMemoryChunk *&prev() { return prev_; }

protected:
    ChainedMemoryChunk *next_;
    ChainedMemoryChunk *prev_;
};

class ChunkAllocatorBase {
public:
    ChunkAllocatorBase() : use_bytes_(0) {}
    virtual ~ChunkAllocatorBase() {}

public:
    void* Allocate(size_t num_bytes) {
        void* p = DoAllocate(num_bytes);
        use_bytes_ += num_bytes;
        return p;
    }
    void Deallocate(void* const addr, size_t num_bytes) {
        DoDeallocate(addr, num_bytes);
        use_bytes_ -= num_bytes;
    }
    size_t UsedBytes() const {
        return use_bytes_;
    }

private:
    virtual void* DoAllocate(size_t num_bytes) = 0;
    virtual void DoDeallocate(void* const addr, size_t num_bytes) = 0;

protected:
    size_t use_bytes_;
};

class MMapAllocator : public ChunkAllocatorBase {
public:
    MMapAllocator() = default;
    ~MMapAllocator() = default;

    void* DoAllocate(size_t n_bytes) override {
        return Mmap(n_bytes);
    }

    void DoDeallocate(void* const addr, size_t num_bytes) override {
        Munmap(addr, num_bytes);
    }

    static void* Mmap(size_t n_bytes) {
        void* addr = ::mmap(0, n_bytes, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
        if (unlikely(addr == MAP_FAILED)) {
            return nullptr;
        }

        madvise(addr, n_bytes, MADV_DONTDUMP);
        return addr;
    }

    static void Munmap(void* const addr, size_t n_bytes) {
        char* const base = (char* const)addr;
        for (size_t offset = 0; offset < n_bytes; offset += MUNMAP_SLICE_SIZE) {
            size_t actual_len = std::min(n_bytes - offset, MUNMAP_SLICE_SIZE);
            munmap(base + offset, actual_len);
            usleep(0);
        }
    }

private:
    static constexpr size_t MUNMAP_SLICE_SIZE = 1024 * 1024;
};

class SimpleAllocator : public ChunkAllocatorBase {
public:
    SimpleAllocator() = default;
    virtual ~SimpleAllocator() = default;

public:
    void* DoAllocate(size_t n_bytes) override {
        return static_cast<void*>(new (std::nothrow) char[n_bytes]);
    }
    void DoDeallocate(void* const addr, size_t n_bytes) override {
        delete[](char *) addr;
    }
};

class AllocatePolicy {
public:
    AllocatePolicy() {}
    virtual ~AllocatePolicy() {}

public:
    /**
     * alloc a chunk
     * @param n_bytes number of bytes to allocate
     * @return allocated chunk
     */
    virtual MemoryChunk* Allocate(size_t n_bytes) = 0;

    /**
     * release allocated chunks
     * @return total size of chunks
     */
    virtual size_t Release() = 0;

    /**
     * Reset memory to init state
     */
    virtual size_t Reset() = 0;

    /**
     * clear trunk data
     */
    virtual void Clear() = 0;

    /**
     * Get used chunk size
     */
    virtual size_t UsedBytes() const = 0;

    /**
     * Get total chunk size
     */
    virtual size_t TotalBytes() const = 0;

    /**
     * Return chunk size
     */
    virtual size_t ChunkSize() const = 0;

    /**
     * Return available chunk size, it is always
     * equal to GetChunkSize() - sizeof(ChainedMemoryChunk)
     */
    virtual size_t AvailableChunkSize() const = 0;

    /*
     * check whether pointer in chunk.
     */
    virtual bool IsInChunk(const void *ptr) const = 0;
};

class SimpleAllocatePolicy : public AllocatePolicy {
public:
    SimpleAllocatePolicy(ChunkAllocatorBase* allocator, size_t chunk_size, bool own_allocator = false);
    SimpleAllocatePolicy(size_t chunk_size);
    virtual ~SimpleAllocatePolicy();

    /**
     * alloc a chunk
     * @param n_bytes number of bytes to allocate
     * @return allocated chunk
     */
    MemoryChunk* Allocate(size_t n_bytes) override;

    /**
     * release allocated chunks
     * @return total size of chunks
     */
    size_t Release() override;

    /**
     * Reset memory to init state
     */
    size_t Reset() override;

    /**
     * clear trunk data
     */
    void Clear() override;

    /**
     * Get used chunk size
     */
    size_t UsedBytes() const override {
        return use_bytes_;
    }

    /**
     * Get total chunk size
     */
    size_t TotalBytes() const override {
        return total_bytes_;
    }

    /**
     * Return chunk size
     */
    size_t ChunkSize() const override {
        return chunk_size_;
    }

    /**
     * Return available chunk size, it is always
     * equal to GetChunkSize() - sizeof(ChainedMemoryChunk)
     */
    size_t AvailableChunkSize() const override {
        return chunk_size_ - sizeof(ChainedMemoryChunk);
    }

    /*
     * check whether pointer in pool.
     */
    bool IsInChunk(const void* ptr) const override;

protected:
    ChunkAllocatorBase* allocator_ = nullptr;
    bool own_allocator_ = false;

    ChainedMemoryChunk* chunk_header_ = nullptr;
    ChainedMemoryChunk* current_chunk_ = nullptr;
    size_t chunk_size_ = 0;
    size_t use_bytes_ = 0;
    size_t total_bytes_ = 0;
};

class PoolBase {
public:
    virtual ~PoolBase() {}

    virtual void* Allocate(size_t num_bytes) = 0;
    virtual void Deallocate(void* ptr, size_t size) = 0;
    virtual void Release() = 0;
    virtual size_t Reset() = 0;
    virtual bool IsInPool(const void* ptr) const = 0;
};

class SimplePool : public PoolBase {
public:
    SimplePool() = default;
    virtual ~SimplePool() = default;

public:
    void* Allocate(size_t num_bytes) override {
        if (num_bytes == 0) {
            return nullptr;
        }

        std::set_new_handler(0);
        void* p = ::operator new(num_bytes);
        if (p) {
            use_bytes_ += num_bytes;
            if (use_bytes_ > peak_of_use_bytes_) {
                peak_of_use_bytes_ = use_bytes_;
            }
        }
        return p;
    }
    void Deallocate(void* ptr, size_t size) override {
        ::operator delete(ptr);
        use_bytes_ -= size;
    }
    void Release() override {
        use_bytes_ = 0;
    }
    size_t Reset() override {
        size_t ret = use_bytes_;
        use_bytes_ = 0;
        return ret;
    }
    bool IsInPool(const void* ptr) const override {
        return false;
    }

    size_t GetUsedBytes() const { return use_bytes_; }
    size_t GetPeakOfUsedBytes() const { return peak_of_use_bytes_; }

private:
    SimplePool(const SimplePool&);
    SimplePool& operator=(const SimplePool&);

    size_t use_bytes_ = 0;
    size_t peak_of_use_bytes_ = 0;
};

class Pool : public PoolBase, public std::enable_shared_from_this<Pool> {
public:
    static const size_t DEFAULT_CHUNK_SIZE = 10 * 1024 * 1024; // 10MB
    static const size_t DEFAULT_ALIGN_SIZE = sizeof(char*);

    Pool(AllocatePolicy* allocate_policy, size_t align_size = DEFAULT_ALIGN_SIZE);
    Pool(ChunkAllocatorBase* allocator, size_t chunk_size, size_t align_size = DEFAULT_ALIGN_SIZE);
    Pool(size_t chunk_size = DEFAULT_CHUNK_SIZE, size_t align_size = DEFAULT_ALIGN_SIZE);
    virtual ~Pool() {
        mem_chunk_ = nullptr;
        delete alloc_policy_;
        alloc_policy_ = nullptr;
    }

    /**
     * Allocate memory from Pool
     * @param n_bytes number of bytes need to allocate
     * @return address of allocated memory
     */
    virtual void* Allocate(size_t n_bytes) override {
        ScopedSpinLock lock(mutex_);
        return AllocateUnsafe(n_bytes);
    }

    virtual void* Allocate(size_t n_bytes, size_t alignment) {
        ScopedSpinLock lock(mutex_);
        return AllocateUnsafe(n_bytes, alignment);
    }

    void* AllocateUnsafe(size_t n_bytes);

    void* AllocateUnsafe(size_t n_bytes, size_t alignment);

    /**
     * Free memory, dummy function
     */
    void Deallocate(void* ptr, size_t size) override {}

    /**
     * release all allocated memory in Pool
     */
    void Release() override {
        ScopedSpinLock lock(mutex_);
        alloc_policy_->Release();
        alloc_size_ = 0;
        mem_chunk_ = &(Pool::dummy_chunk_);
    }

    /**
     * Reset memory to init state
     */
    size_t Reset() override {
        ScopedSpinLock lock(mutex_);
        return ResetUnsafe();
    }
    size_t ResetUnsafe() {
        size_t total_size = alloc_policy_->Reset();
        mem_chunk_ = &(Pool::dummy_chunk_);
        alloc_size_ = 0;
        return total_size;
    }

    /**
     * clear pool data and reset
     */
    void Clear() {
        ScopedSpinLock lock(mutex_);
        alloc_policy_->Clear();
    }

    void* AllocateAlign(size_t n_bytes, size_t align_size) {
        n_bytes += align_size;
        return (void*)AlignBytes((size_t)Allocate(n_bytes), align_size);
    }
    /**
     * get alignment size
     * @return align_size_ alignment size
     */
    size_t AlignSize() const {
        return alloc_size_;
    }

    /**
     * Get used size in bytes
     */
    virtual size_t UsedBytes() {
        ScopedSpinLock lock(mutex_);
        size_t alloc_bytes = alloc_policy_->UsedBytes();
        alloc_bytes -= mem_chunk_->FreeSize();
        return alloc_bytes;
    }

    /**
     * Get total size in bytes
     */
    virtual size_t TotalBytes() {
        ScopedSpinLock lock(mutex_);
        return alloc_policy_->TotalBytes();
    }

    /**
     *Get allocated size from pool
     */
    size_t AllocatedSize() {
        ScopedSpinLock lock(mutex_);
        return alloc_size_;
    }

    size_t AvailableChunkSize() {
        ScopedSpinLock lock(mutex_);
        return alloc_policy_->AvailableChunkSize();
    }

    /*
     * Check whether pointer in pool
     */
    bool IsInPool(const void *ptr) const override {
        return alloc_policy_->IsInChunk(ptr);
    }

    static size_t AlignBytes(size_t n_bytes, size_t align_size_) {
        return ((n_bytes + align_size_ - 1) & ~(align_size_ - 1));
    }

    static std::shared_ptr<Pool> MaybeShared(Pool* pool);

protected:
    SpinLock mutex_;
    size_t align_size_;
    MemoryChunk* mem_chunk_;
    size_t alloc_size_;

    AllocatePolicy* alloc_policy_;
    static MemoryChunk dummy_chunk_;

private:
    Pool(const Pool &);
    void operator=(const Pool &);
};


template <typename T>
class pool_allocator {
public:
    typedef size_t size_type;
    typedef ptrdiff_t difference_type;
    typedef T* pointer;
    typedef const T* const_pointer;
    typedef T& reference;
    typedef const T& const_reference;
    typedef T value_type;

    template <typename _Tp1>
    struct rebind {
        typedef pool_allocator<_Tp1> other;
    };

    pool_allocator(Pool *pool) throw() : _pool(pool) {}

    pool_allocator(const pool_allocator &other) throw() : _pool(other._pool) {}

    template <typename _Tp1>
    pool_allocator(const pool_allocator<_Tp1> &other) throw() : _pool(other._pool) {}

    ~pool_allocator() throw() {}

    pointer address(reference __x) const { return &__x; }

    const_pointer address(const_reference __x) const { return &__x; }

    // NB: __n is permitted to be 0.  The C++ standard says nothing
    // about what the return value is when __n == 0.
    pointer allocate(size_type __n, const void * = 0) {
        if (__builtin_expect(!_pool, false))
            std::__throw_bad_alloc();
        return static_cast<T*>(_pool->Allocate(__n * sizeof(T)));
    }

    void deallocate(pointer __p, size_type __n) {
        if (_pool) {
            _pool->Deallocate((void *)__p, __n * sizeof(T));
        }
    }

    size_type max_size() const throw() { return size_t(-1) / sizeof(T); }

    void construct(pointer __p, const T &__val) { ::new (__p) T(__val); }

    void construct(pointer __p, T &&__val) { ::new (__p) T(std::move(__val)); }

    void destroy(pointer __p) { __p->~T(); }

public:
    Pool *_pool;
};

class RecyclePool : public Pool {
public:
public:
    RecyclePool(ChunkAllocatorBase* allocator, size_t chunkSize, size_t alignSize = 1);
    RecyclePool(size_t chunkSize, size_t alignSize = 1);
    ~RecyclePool();

private:
    RecyclePool(const RecyclePool &);
    void operator=(const RecyclePool &);

public:
    /**
     * num_bytes >= 8
     */
    virtual void* Allocate(size_t num_bytes) override;

    /*
     * only less than 64k memory block can be deallocated
     */
    virtual void Deallocate(void* ptr, size_t size) override;
    virtual void Release() override;
    virtual size_t Reset() override;

    size_t GetFreeSize() const;

public:
    // for test
    std::unordered_map<size_t, void*>& GetFreeList();

private:
    static void*& NextOf(void* ptr);

private:
    size_t free_size_;
    std::unordered_map<size_t, void *> free_list_;
};

inline void*& RecyclePool::NextOf(void* ptr) { return *(static_cast<void **>(ptr)); }

inline void* RecyclePool::Allocate(size_t num_bytes) {
    size_t alloc_size = AlignBytes(num_bytes, align_size_);
    if (alloc_size < 8) {
        alloc_size = 8;
    }
    std::unordered_map<size_t, void*>::iterator it = free_list_.find(alloc_size);
    if (it != free_list_.end()) {
        void* ptr = it->second;
        if (ptr != NULL) {
            free_list_[alloc_size] = NextOf(ptr);
            free_size_ -= alloc_size;
            alloc_size_ += alloc_size;
            return ptr;
        }
    }
    return Pool::Allocate(alloc_size);
}

inline void RecyclePool::Deallocate(void* ptr, size_t size) {
    if (!ptr || size == 0) {
        return;
    }
    size_t alloc_size = AlignBytes(size, align_size_);
    if (alloc_size < 8) {
        alloc_size = 8;
    }

    std::unordered_map<size_t, void*>::iterator it = free_list_.find(alloc_size);
    if (it == free_list_.end()) {
        free_list_[alloc_size] = NULL;
    }
    NextOf(ptr) = free_list_[alloc_size];
    free_list_[alloc_size] = ptr;
    alloc_size_ -= alloc_size;
    free_size_ += alloc_size;
}

inline void RecyclePool::Release() {
    free_size_ = 0;
    free_list_.clear();

    Pool::Release();
}

inline size_t RecyclePool::Reset() {
    free_size_ = 0;
    free_list_.clear();

    return Pool::Reset();
}

inline size_t RecyclePool::GetFreeSize() const { return free_size_; }
inline std::unordered_map<size_t, void*>& RecyclePool::GetFreeList() { return free_list_; }

///////////////////////////////////////////////////////////////////////////////////////

template <typename T>
inline bool operator==(const pool_allocator<T> &lhs, const pool_allocator<T> &rhs) {
    return lhs._pool == rhs._pool;
}

template <typename T>
inline bool operator!=(const pool_allocator<T> &lhs, const pool_allocator<T> &rhs) {
    return !(lhs == rhs);
}

#define IE_POOL_NEW_CLASS(pool, type, args...) (new ((pool)->Allocate(sizeof(type))) type(args))
#define IE_POOL_COMPATIBLE_NEW_CLASS(pool, type, args...)                                                              \
    ((pool) ? IE_POOL_NEW_CLASS(pool, type, args) : new type(args))

template <typename T>
inline void IE_POOL_DELETE_CLASS(Pool* pool, T* value)
{
    if (value) {
        value->~T();
        pool->Deallocate((void*)value, sizeof(T));
    }
}
template <typename T>
inline void IE_POOL_COMPATIBLE_DELETE_CLASS(Pool* pool, T*& value)
{
    (pool) ? IE_POOL_DELETE_CLASS(pool, value) : (delete (value));
    value = nullptr;
}

///////////////////////////////////////////////////////////////////////////////////////

#define IE_POOL_NEW_VECTOR(pool, type, num) ((type*)((pool)->Allocate(sizeof(type) * num)))
template <typename T>
inline void IE_POOL_DELETE_VECTOR(Pool* pool, T* value, size_t num)
{
    if (value) {
        pool->Deallocate((void*)value, sizeof(T) * num);
    }
}

#define IE_POOL_COMPATIBLE_NEW_VECTOR(pool, type, num) (pool) ? IE_POOL_NEW_VECTOR(pool, type, num) : new type[num]
#define IE_POOL_COMPATIBLE_DELETE_VECTOR(pool, value, num)                                                             \
    if (!pool) {                                                                                                       \
        delete[] value;                                                                                                \
    } else {                                                                                                           \
        indexlib::IE_POOL_DELETE_VECTOR(pool, value, num);                                                             \
    }                                                                                                                  \
    value = nullptr

#define POOL_NEW_CLASS(pool, type, args...) (new (pool->Allocate(sizeof(type))) type(args))

template <typename T>
inline void POOL_DELETE_CLASS(T *value) {
    if (value) {
        value->~T();
    }
}

}
}
